#!/usr/bin/python
# Counts kmers with no mismatch.
import sys
import numpy as np

# parameters.
in_file = sys.argv[1]
out_file = sys.argv[2]
#num_contigs = int(sys.argv[3])

k = 25
step = 5
edge = 3000
cutoff = 4

# declarations.
names_dt = np.dtype([\
		('kmer', np.str_, k),\
		('name', np.str_, 200),\
		('start', np.int),\
		])
		

# read in contigs.
def read_fasta(file_path):
	# read in files.
	fin = open(file_path, "rb")
	lines = fin.readlines()
	fin.close()
	
	# build dictionary.
	seq_dat = {}
	first = True
	for line in lines:
		# check if we start new entry.
		if line[0] == ">":
			# build kmer 1 mismatch.
			if first != True:
				# save to dictionary.
				seq_dat[head] = seq
		
			else:
				# clear first.
				first = False
				
			# clear.
			head = line.strip().replace(">","")
			seq = ""
			continue
		
		
		# get seq
		seq += line.strip()
	# save last tp dictionary.
	seq_dat[head] = seq	
	
	# cleanup and return.
	del lines
	return seq_dat

def get_edges(seq, head, sidx):

	# cut into start kmers
	max = edge - k
	if max > len(seq):
		max = len(seq) - k
	for i in range(0, max, step):
		# get kmer.
		key = seq[i:i+k]
	
		# increment kmer count.
		if key not in db:
			db[key] = 0
		db[key] += 1
		
		# add info.
		kmers[sidx][0] = key
		kmers[sidx][1] = head
		kmers[sidx][2] = i
		sidx += 1
		
			
	# cut into stop kmers
	max = len(seq) - edge
	if max < 0:
		max = 0
	for i in range(max, len(seq)-k, step):
		# get kmer.
		key = seq[i:i+k]
		
		# increment kmer count.
		if key not in db:
			db[key] = 0
		db[key] += 1
		
		# add info.
		kmers[sidx][0] = key
		kmers[sidx][1] = head
		kmers[sidx][2] = i
		sidx += 1
			
			
	return sidx

########
# read contigs.
print "reading contigs."
seq_data = read_fasta(in_file)
num_contigs = len(seq_data)

# Initialize data strucutres.
size = 2 * num_contigs * int(float(edge)/float(step))
print "initializing data structure.", size
kmers = np.zeros(size,dtype=names_dt)
db = {}
	
# process file.
print "processing kmers out."
sidx = 0
seq_cnt = 0
for ctg in seq_data:
	# Process kmers.
	seq = seq_data[ctg]
	sidx = get_edges(seq, ctg, sidx)

	# track progress.
	seq_cnt += 1
	if seq_cnt % 1000 == 0: print seq_cnt
	

# print out contig, start, stop of bad positions.
print "printing bad parts."
fout = open(out_file, "wb")

for kmer in kmers:
	# ensure consistant.
	if kmer[0] not in db: continue
	
	# Check if good.
	if db[kmer[0]] < cutoff:
		continue
		
	# print indices.
	fout.write("%i\t%s\t%i\t%i\n" % (db[kmer[0]], kmer[1], kmer[2], kmer[2] + k))
fout.close()






